# This code is imported from the replication files of Knox et al (2019). 

'%.%' <- function(x,y) paste(x,y,sep='')

## data manipulation
library(reshape2)
library(plyr)
library(MASS)    # for mvrnorm()
library(gtools)  # for rdirichlet()
library(mlogit)

d.raw <- read.csv('Spring2014omnibus_raw.csv', stringsAsFactors = FALSE)

source('bounds_frechet_function.R')
source('lll.R')

draw_start <- 1
draw_stop <- 100
choice_divergences  <- 1/3
outcome_divergences <- 0

#####################
## preprocess data ##
#####################

## media assignment variables: convert NA/1 to 0/1
d.raw$fox[is.na(d.raw$fox)] <- 0
d.raw$msnbc[is.na(d.raw$msnbc)] <- 0
d.raw$entertainment[is.na(d.raw$entertainment)] <- 0
d.raw$med_choice[is.na(d.raw$med_choice)] <- 0  # 0 if no choice given, else 1-4

# d.raw_1 <- subset(d.raw, forcedchoice == 1)
# d.raw_0 <- subset(d.raw, forcedchoice == 0)
# d.raw_0 <- subset(d.raw, med_choice >= 1)
# d.raw <- rbind(d.raw_1, d.raw_0)

## construct 'stated preference' variable
d.raw$med_pref[which(d.raw$med_pref == 4)] <- 3
d.raw$med_pref[which(d.raw$med_pref == 0)] <- NA

## construct 'actual choice' variable
##   (pool both entertainment options in free-choice arm,
##    already pooled in forced choice arm)
d.raw$med_choice[which(d.raw$med_choice == 4)] <- 3
d.raw$med_choice[which(d.raw$med_choice == 0)] <- NA

## construct 'received treatment' assignment variable
##   (A=c if D=0, A~cat(1/3,1/3,1/3) if D=1)
##   (fox = 1, msnbc = 2, entertainment = 3)
d.raw$med_assign <-
  ifelse(d.raw$forcedchoice == 1,
         d.raw$fox + 2*d.raw$msnbc + 3*d.raw$entertainment,
         d.raw$med_choice
  )

## construct party id variable
d.raw$pid <- NA
d.raw$pid[d.raw$party1 == 1] <- -1  # considers self a democrat
d.raw$pid[d.raw$party1 == 2] <- 1   # considers self a republican
d.raw$pid[d.raw$party4 == 1] <- 1   # neither but closer to rep
d.raw$pid[d.raw$party4 == 2] <- -1  # neither but closer to rep
d.raw$pid[d.raw$party4 == 3] <- 0   # closer to neither rep or dem

## construct pro/counter-attitudinal media variable
##   pro-attitudinal = 1      (fox for reps, or msnbc for dems)
##   counter-attitudinal = 2  (msnbc for reps, or fox for dems)
##   entertainment = 3        (for both parties)
med_labels <- c('pro-attitudinal','counter-attitudinal','entertainment')
d.raw <- d.raw[-which(d.raw$pid == 0 | is.na(d.raw$pid)),] # drop independents
d.raw$med_pref.att <-
  ifelse((d.raw$med_pref == 1 & d.raw$pid == 1) | (d.raw$med_pref == 2 & d.raw$pid == -1), 1,
         ifelse((d.raw$med_pref == 2 & d.raw$pid == 1) | (d.raw$med_pref == 1 & d.raw$pid == -1), 2, 3)
  )
d.raw$med_choice.att <-
  ifelse((d.raw$med_choice == 1 & d.raw$pid == 1) | (d.raw$med_choice == 2 & d.raw$pid == -1), 1,
         ifelse((d.raw$med_choice == 2 & d.raw$pid == 1) | (d.raw$med_choice == 1 & d.raw$pid == -1), 2, 3)
  )
d.raw$med_assign.att <-
  ifelse((d.raw$med_assign == 1 & d.raw$pid == 1) | (d.raw$med_assign == 2 & d.raw$pid == -1), 1,
         ifelse((d.raw$med_assign == 2 & d.raw$pid == 1) | (d.raw$med_assign == 1 & d.raw$pid == -1), 2, 3)
  )

## construct media sentiment index and rescale to [0,1]
##   these are eight questions on perception of media  (7-point scales)
##   name of variable is lower end of scale
##   first make direction consistent between questions
d.raw$perceive.index <- rowSums(cbind(d.raw$fair_4,
                                      d.raw$friendly_4,
                                      d.raw$good_4,
                                      d.raw$quarrel_4 * -1 + 8,
                                      d.raw$balanced_4,
                                      d.raw$oneside_4* -1 + 8,
                                      d.raw$american_4,
                                      d.raw$accurate_4
))
##   then flip so larger numbers mean 'media is good'
##   for now we'll keep integer values for simulation
##   instead of rescaling like in main analysis (convert later)
d.raw$perceive.index <- -(d.raw$perceive.index - 8) + 48

## make clean data for analysis with [0,1] outcome:
##   media sentiment index
d <- data.frame(S = d.raw$med_pref.att,
                C = d.raw$med_choice.att,
                A = d.raw$med_assign.att,
                D = d.raw$forcedchoice
)
d$Y <- d.raw$perceive.index
drop <- which(is.na(d$S) |
                is.na(d$D) |
                is.na(d$A) |
                is.na(d$Y) |
                (is.na(d$C) & d$D == 0)
)
if (length(drop) > 0){
  d <- d[-drop,]
}
rm(drop)


## prepare to estimate effect of changing between these treatment values
J <- 3
effects.mat <- matrix(FALSE, J, J)
dimnames(effects.mat) <- list(ctrl = NULL,
                              treat = NULL)
effects.mat[2, 1] <- TRUE
effects.mat[3, 1] <- TRUE
effects.mat[3, 2] <- TRUE

########################
## summary statistics ##
########################

# sum(1{S=s,C=c})
n_sc = table(S=d$S,C=d$C)
rownames(n_sc) = "S." %.% 1:J
colnames(n_sc) = "C." %.% 1:J

# sum(1{S=s,A=a})
n_sa_pool = table(S=d$S,A=d$A)
rownames(n_sa_pool) = "S." %.% 1:J
colnames(n_sa_pool) = "A." %.% 1:J

# sum(1{S=s,A=a,D=1})
n_sa_force = table(S=d$S[d$D==1],A=d$A[d$D==1])
rownames(n_sa_force) = "S." %.% 1:J
colnames(n_sa_force) = "A." %.% 1:J

# Pr(S=s)
p_s = table(d$S) / nrow(d)

# Pr(C=c|S=s)
p_c.s = t(sapply(1:J,function(s){
  out = table(d[d$S==s & d$D==0,"C"]) / length(d[d$S==s & d$D==0,"C"])
  names(out) = "C." %.% names(out)
  return(out)
}))
rownames(p_c.s) = "S." %.% 1:J

# Pr(S=s,C=c)
p_sc_free = n_sc / sum(n_sc)

p_sc_pool = p_s %o% rep(1,J) * p_c.s
dimnames(p_sc_pool) = dimnames(p_sc_free)

p_c = colSums(p_sc_pool)

p_s.c = p_sc_pool / (rep(1,J) %o% p_c)



### observed outcomes ###

# from forced condition: E[Y|S=s,A=a,D=1]
y_sa = t(sapply(1:J,function(s){
  out = sapply(1:J,function(a){
    mean(d$Y[which(d$S==s & d$A==a & d$D==1)],na.rm=T)
  })
  names(out) = "A." %.% 1:J
  return(out)
}))
rownames(y_sa) = "S." %.% 1:J

# from free condition: E[Y|S=s,C=c,A=c,D=0]
pi_scc = t(sapply(1:J,function(s){
  out = sapply(1:J,function(c){
    mean(d$Y[which(d$S==s & d$C==c & d$D==0)],na.rm=T)
  })
  names(out) = "C." %.% 1:J
  return(out)
}))
rownames(pi_scc) = "S." %.% 1:J

##################################################
# create DGP that is consistent with pilot study #
##################################################

sc.comb = as.matrix(expand.grid(c=1:J,s=1:J)[,2:1])
sa.comb = sc.comb; colnames(sa.comb) = c("s","a")
sca.comb = as.matrix(expand.grid(a=1:J,c=1:J,s=1:J)[,3:1])

# This is where we generate DGP 
make_sim_data_placebo = function(n, choice_divergence = 0, outcome_divergence = 0,
                                 placebo = 3, vio = 0){
  
  u = choice_divergence
  v = outcome_divergence
  
  ## simulation parameter u in [0,1] determines how informative the
  ##   stated choices:
  ##     u=0 means that stated choices are as informative as in real data
  ##     u=1 means that stated choices are totally uninformative and
  ##         actual choices are uniformly distributed regardless of statement
  observed_sc = p_c.s  # sum over rows equal to 1.
  uninformative_sc = matrix(1/J, J, J, dimnames = dimnames(p_c.s))
  p_c.s_simpop = u * uninformative_sc + (1-u) * observed_sc
  
  ## fY_sca[s,c,a,] is vector of Pr( Y(a)=y | S=s, C=c )
  ##   where y = 0, ..., 48
  fY_sca_simpop = array(NA, dim=c(J,J,J, 49))
  dimnames(fY_sca_simpop) = list(S = "S." %.% 1:J,
                                 C = "C." %.% 1:J,
                                 A = "A." %.% 1:J,
                                 p_Y = "Y." %.% 0:48
  )
  
  # plug in data from free-choice condition
  for (row in 1:nrow(sc.comb)){
    s = sc.comb[row,1]
    c = sc.comb[row,2]
    out = table(d$Y[d$S==s & d$C==c & d$A==c & d$D==0])
    names(out) = "Y." %.% names(out)
    out = out["Y." %.% 0:48]
    out[is.na(out)] = 0
    fY_sca_simpop[s,c,c,] = out / sum(out)
  }
  
  ## fY_sa[s,a,] is vector of Pr( Y(a)=y | S=s )
  ##   where y = 0, ..., 48
  fY_sa = array(NA, dim=c(J,J,49))
  dimnames(fY_sa) = list(S = "S." %.% 1:J,
                         A = "A." %.% 1:J,
                         p_Y = "Y." %.% 0:48
  )
  # plug in data from forced-choice condition
  for (row in 1:nrow(sa.comb)){
    s = sa.comb[row,1]
    a = sa.comb[row,2]
    out = table(d$Y[d$S==s & d$A==a & d$D==1])
    names(out) = "Y." %.% names(out)
    out = out["Y." %.% 0:48]
    out[is.na(out)] = 0
    fY_sa[s,a,] = out / sum(out)
  }
  
  ## Observed forced-choice f( Y(a) | S=s ) is mixture of 3 components,
  ##   f( Y(a) | S=s, C=c),
  ##   where f( Y(a) | S=s, C=a ) is copied from free choice condition data.
  ## Simulation parameter outcome_divergence (v in [0, 1]) determines
  ##   how divergent the remaining two components are:
  ##     v=0 means that they are identical in distribution and
  ##     v=1 means that max of one <= min of other
  
  for (row in 1:nrow(sa.comb)){
    s = sa.comb[row,1]
    a = sa.comb[row,2]
    
    forced_mixture = fY_sa[s,a,]
    free_component = fY_sca_simpop[s,a,a,] * p_c.s_simpop[s,a]
    
    ## fix issues arising from the fact that our data is a finite sample
    unallocated = pmax(forced_mixture - free_component, 0)
    unallocated = unallocated / sum(unallocated) * (1 - p_c.s_simpop[s,a])
    
    c1 = c(2, 3, 1)[c(2, 3, 1) != a][1]
    c2 = (1:J)[!1:J %in% c(a, c1)]
    ## most balanced allocation:
    ##   divide unexplained distr evenly among choice groups by their size
    most.balanced.allocation =
      unallocated * (p_c.s_simpop[s,c1] / (1 - p_c.s_simpop[s,a]))
    ## most extreme allocation:
    ##   place portion of distr with lowest values into first choice group
    ##   (up to its size) and highest into other choice group
    unallocated.cumulative = cumsum(unallocated)
    most.extreme.allocation = unallocated
    most.extreme.allocation[unallocated.cumulative > p_c.s_simpop[s,c1]] = 0
    split.mass.at.this.y = which(unallocated.cumulative > p_c.s_simpop[s,c1])[1]
    most.extreme.allocation[split.mass.at.this.y] =
      p_c.s_simpop[s,c1] - c(0, unallocated.cumulative)[split.mass.at.this.y]
    
    ## select allocation by smoothing between extreme and balanced
    allocation = v * most.extreme.allocation + (1-v) * most.balanced.allocation
    fY_sca_simpop[s,c1,a,] = allocation / p_c.s_simpop[s,c1]
    fY_sca_simpop[s,c2,a,] = (unallocated - allocation) / p_c.s_simpop[s,c2]
    
  }
  
  ## stated & true prefs
  p_sc_simpop = p_s %o% rep(1,J) * p_c.s_simpop
  p_c_simpop = colSums(p_sc_simpop)
  p_s.c_simpop = p_sc_simpop / (rep(1,J) %o% p_c_simpop)
  
  ## Pr(Y(a)<=y | S=s)
  F_Y_sa = t(sapply(1:J,function(s){
    sapply(1:J,function(a){
      stepfun(
        x = 0:48,
        y = c(cumsum(p_c.s_simpop[s,] %*% fY_sca_simpop[s,,a,]),
              1
        )
      )
    })
  }))
  dimnames(F_Y_sa) = list(S=NULL,A=NULL)
  
  ## Pr(Y(c)<=y | S=s,C=c)
  F_Y_scc = t(sapply(1:J,function(s){
    sapply(1:J,function(c){
      stepfun(
        x = 0:48,
        y = c(cumsum(fY_sca_simpop[s,c,c,]), 1)
      )
    })
  }))
  dimnames(F_Y_scc) = list(S=NULL,C=NULL)
  
  ## Pr(Y(a)<=y | S=s, C!=c)
  F_diff_Y_sa =  t(sapply(1:J,function(s){
    sapply(1:J,function(a){
      ## evaluate 's' and 'a' in local context (single iteration of sapply)
      ## Otherwise, they are evaluated in the parent environment, of 'sapply' itself
      ##   meaning that s=J and a=J (ending values in the parent)
      env_sa = new.env()
      assign("s",s,envir=env_sa) # save current value of 's' into env_sa, where it will remain unchanged by future iterations
      assign("a",a,envir=env_sa) # save current value of 'a' into env_sa, where it will remain unchanged by future iterations
      f_sa = function(y){
        out = (F_Y_sa[[s,a]](y) - F_Y_scc[[s,a]](y) * p_c.s_simpop[s,a]) / (1 - p_c.s_simpop[s,a])
        return(out)
      }
      environment(f_sa) = env_sa
      attr(f_sa,"knots") = sort(unique(c(0,48,knots(F_Y_scc[[s,a]]),knots(F_Y_sa[[s,a]]))))
      attr(f_sa,"values") = sapply(attr(f_sa,"knots"),f_sa)
      return(f_sa)
    })
  }))
  dimnames(F_diff_Y_sa) = list(S=NULL,A=NULL)
  
  ## sampling a new dataset
  
  n_sc_true = matrix(rmultinom(1,n,as.numeric(p_sc_simpop)),J,J,
                     dimnames=list(S="S." %.% 1:J,C="C." %.% 1:J))
  
  data = do.call(rbind,sapply(1:J,function(c){
    do.call(rbind,sapply(1:J,function(s){
      data.frame(S=rep(s,n_sc_true[s,c]),
                 C=rep(c,n_sc_true[s,c]))
    },simplify=FALSE))
  },simplify=FALSE))
  
  # ##########################
  # Observe Data 
  # ##########################
  # This is the only difference from before 
  
  data$D = sample(rep(0:1, each=n/2))
  
  # This is the only difference from before
  # Free-Choice
  # data$pl <- rep(0, nrow(data))
  # data$pl[data$D == 0 & data$C != placebo] <- 
  #  sample(x = c(0, 1), size = sum(data$D == 0 & data$C != placebo), replace = TRUE)
  # data$A[data$D == 0 & data$pl == 0] = data$C[data$D==0 & data$pl == 0]
  #
  # data$A[data$D == 0 & data$pl == 1] = 4 # Using 4 to denote placebo 
  
  # Free-Choice
  data$pl <- rep(0, nrow(data))
  data$pl[data$D == 0 & data$C == 1] <- rbinom(n = sum(data$D==0 & data$C == 1), size = 1, prob = 0.5)
  data$pl[data$D == 0 & data$C == 2] <- rbinom(n = sum(data$D==0 & data$C == 2), size = 1, prob = 0.5)
  data$pl[data$D == 0 & data$C == 3] <- rep(0, sum(data$D==0 & data$C == 3))
  # data$pl[data$D == 0 & data$C != placebo] <- 
  #   sample(x = c(0, 1), size = sum(data$D == 0 & data$C != placebo), replace = TRUE)
  
  data$A[data$D == 0 & data$pl == 0] = data$C[data$D==0 & data$pl == 0]
  data$A[data$D == 0 & data$pl == 1] = 4 # Using 4 to denote placebo 
  
  # Forced Choice
  data$A[data$D==1] = sample(1:J, n/2, replace=T)
  
  data$Y = NA
  for (s in 1:J){
    for (c in 1:J){
      for (a in 1:J){
        # Placebo = 0
        data$Y[data$S==s & data$C==c & data$A==a & data$pl == 0] =
          sample(0:48, size=sum(data$S==s & data$C==c & data$A==a & data$pl == 0),
                 prob = fY_sca_simpop[s,c,a,],
                 replace=T)
      }
    }
  }
  # Placebo 
  for (s in 1:J){
    for (c in 1:J){
      # Placebo = 1
      # If the assumption is correct, we observe potential outcomes under 3
      tilt <- rbinom(n = 1, size = 1, prob = vio)
      base_value <- sample(0:48, size=sum(data$S==s & data$C==c & data$pl == 1),
                           prob = fY_sca_simpop[s,c,placebo,], # fY_sca_simpop[s,c,a,]
                           replace=T)
      if(tilt == 1){
        data$Y[data$S==s & data$C==c & data$pl == 1] <- base_value/2
      }else if(tilt == 0){
        data$Y[data$S==s & data$C==c & data$pl == 1] <- base_value
      }
    }
  }
  
  data$C[data$D==1] = NA
  
  return(list(data=data,
              fY_sca_simpop = fY_sca_simpop,
              p_sc_simpop = p_sc_simpop,
              p_c_simpop = p_c_simpop,
              p_s.c_simpop = p_s.c_simpop,
              p_c.s_simpop = p_c.s_simpop,
              F_Y_sa = F_Y_sa,
              F_Y_scc = F_Y_scc,
              F_diff_Y_sa = F_diff_Y_sa
  ))
  
}

# ########################
# PICA1
# #########################
make_sim_data = function(n, choice_divergence = 0, outcome_divergence = 0){
  
  u = choice_divergence
  v = outcome_divergence
  
  ## simulation parameter u in [0,1] determines how informative the
  ##   stated choices:
  ##     u=0 means that stated choices are as informative as in real data
  ##     u=1 means that stated choices are totally uninformative and
  ##         actual choices are uniformly distributed regardless of statement
  observed_sc = p_c.s
  uninformative_sc = matrix(1/J, J, J, dimnames = dimnames(p_c.s))
  p_c.s_simpop = u * uninformative_sc + (1-u) * observed_sc
  
  ## fY_sca[s,c,a,] is vector of Pr( Y(a)=y | S=s, C=c )
  ##   where y = 0, ..., 48
  fY_sca_simpop = array(NA, dim=c(J,J,J,49))
  dimnames(fY_sca_simpop) = list(S = "S." %.% 1:J,
                                 C = "C." %.% 1:J,
                                 A = "A." %.% 1:J,
                                 p_Y = "Y." %.% 0:48
  )
  
  # plug in data from free-choice condition
  for (row in 1:nrow(sc.comb)){
    s = sc.comb[row,1]
    c = sc.comb[row,2]
    out = table(d$Y[d$S==s & d$C==c & d$A==c & d$D==0])
    names(out) = "Y." %.% names(out)
    out = out["Y." %.% 0:48]
    out[is.na(out)] = 0
    fY_sca_simpop[s,c,c,] = out / sum(out)
  }
  
  ## fY_sa[s,a,] is vector of Pr( Y(a)=y | S=s )
  ##   where y = 0, ..., 48
  fY_sa = array(NA, dim=c(J,J,49))
  dimnames(fY_sa) = list(S = "S." %.% 1:J,
                         A = "A." %.% 1:J,
                         p_Y = "Y." %.% 0:48
  )
  # plug in data from forced-choice condition
  for (row in 1:nrow(sa.comb)){
    s = sa.comb[row,1]
    a = sa.comb[row,2]
    out = table(d$Y[d$S==s & d$A==a & d$D==1])
    names(out) = "Y." %.% names(out)
    out = out["Y." %.% 0:48]
    out[is.na(out)] = 0
    fY_sa[s,a,] = out / sum(out)
  }
  
  ## Observed forced-choice f( Y(a) | S=s ) is mixture of 3 components,
  ##   f( Y(a) | S=s, C=c),
  ##   where f( Y(a) | S=s, C=a ) is copied from free choice condition data.
  ## Simulation parameter outcome_divergence (v in [0, 1]) determines
  ##   how divergent the remaining two components are:
  ##     v=0 means that they are identical in distribution and
  ##     v=1 means that max of one <= min of other
  
  for (row in 1:nrow(sa.comb)){
    s = sa.comb[row,1]
    a = sa.comb[row,2]
    
    forced_mixture = fY_sa[s,a,]
    free_component = fY_sca_simpop[s,a,a,] * p_c.s_simpop[s,a]
    
    ## fix issues arising from the fact that our data is a finite sample
    unallocated = pmax(forced_mixture - free_component, 0)
    unallocated = unallocated / sum(unallocated) * (1 - p_c.s_simpop[s,a])
    
    c1 = c(2, 3, 1)[c(2, 3, 1) != a][1]
    c2 = (1:J)[!1:J %in% c(a, c1)]
    ## most balanced allocation:
    ##   divide unexplained distr evenly among choice groups by their size
    most.balanced.allocation =
      unallocated * (p_c.s_simpop[s,c1] / (1 - p_c.s_simpop[s,a]))
    ## most extreme allocation:
    ##   place portion of distr with lowest values into first choice group
    ##   (up to its size) and highest into other choice group
    unallocated.cumulative = cumsum(unallocated)
    most.extreme.allocation = unallocated
    most.extreme.allocation[unallocated.cumulative > p_c.s_simpop[s,c1]] = 0
    split.mass.at.this.y = which(unallocated.cumulative > p_c.s_simpop[s,c1])[1]
    most.extreme.allocation[split.mass.at.this.y] =
      p_c.s_simpop[s,c1] - c(0, unallocated.cumulative)[split.mass.at.this.y]
    
    ## select allocation by smoothing between extreme and balanced
    allocation = v * most.extreme.allocation + (1-v) * most.balanced.allocation
    fY_sca_simpop[s,c1,a,] = allocation / p_c.s_simpop[s,c1]
    fY_sca_simpop[s,c2,a,] = (unallocated - allocation) / p_c.s_simpop[s,c2]
    
  }
  
  ## stated & true prefs
  p_sc_simpop = p_s %o% rep(1,J) * p_c.s_simpop
  p_c_simpop = colSums(p_sc_simpop)
  p_s.c_simpop = p_sc_simpop / (rep(1,J) %o% p_c_simpop)
  
  ## Pr(Y(a)<=y | S=s)
  F_Y_sa = t(sapply(1:J,function(s){
    sapply(1:J,function(a){
      stepfun(
        x = 0:48,
        y = c(cumsum(p_c.s_simpop[s,] %*% fY_sca_simpop[s,,a,]),
              1
        )
      )
    })
  }))
  dimnames(F_Y_sa) = list(S=NULL,A=NULL)
  
  ## Pr(Y(c)<=y | S=s,C=c)
  F_Y_scc = t(sapply(1:J,function(s){
    sapply(1:J,function(c){
      stepfun(
        x = 0:48,
        y = c(cumsum(fY_sca_simpop[s,c,c,]), 1)
      )
    })
  }))
  dimnames(F_Y_scc) = list(S=NULL,C=NULL)
  
  ## Pr(Y(a)<=y | S=s, C!=c)
  F_diff_Y_sa =  t(sapply(1:J,function(s){
    sapply(1:J,function(a){
      ## evaluate 's' and 'a' in local context (single iteration of sapply)
      ## Otherwise, they are evaluated in the parent environment, of 'sapply' itself
      ##   meaning that s=J and a=J (ending values in the parent)
      env_sa = new.env()
      assign("s",s,envir=env_sa) # save current value of 's' into env_sa, where it will remain unchanged by future iterations
      assign("a",a,envir=env_sa) # save current value of 'a' into env_sa, where it will remain unchanged by future iterations
      f_sa = function(y){
        out = (F_Y_sa[[s,a]](y) - F_Y_scc[[s,a]](y) * p_c.s_simpop[s,a]) / (1 - p_c.s_simpop[s,a])
        return(out)
      }
      environment(f_sa) = env_sa
      attr(f_sa,"knots") = sort(unique(c(0,48,knots(F_Y_scc[[s,a]]),knots(F_Y_sa[[s,a]]))))
      attr(f_sa,"values") = sapply(attr(f_sa,"knots"),f_sa)
      return(f_sa)
    })
  }))
  dimnames(F_diff_Y_sa) = list(S=NULL,A=NULL)
  
  ## sampling a new dataset
  
  n_sc_true = matrix(rmultinom(1,n,as.numeric(p_sc_simpop)),J,J,
                     dimnames=list(S="S." %.% 1:J,C="C." %.% 1:J))
  
  data = do.call(rbind,sapply(1:J,function(c){
    do.call(rbind,sapply(1:J,function(s){
      data.frame(S=rep(s,n_sc_true[s,c]),
                 C=rep(c,n_sc_true[s,c]))
    },simplify=FALSE))
  },simplify=FALSE))
  
  data$D = sample(rep(0:1,each=n/2))
  
  data$A[data$D==0] = data$C[data$D==0]
  data$A[data$D==1] = sample(1:J,n/2,replace=T)
  
  data$Y = NA
  for (s in 1:J){
    for (c in 1:J){
      for (a in 1:J){
        data$Y[data$S==s & data$C==c & data$A==a] =
          sample(0:48, size=sum(data$S==s & data$C==c & data$A==a),
                 prob=fY_sca_simpop[s,c,a,],
                 replace=T)
      }
    }
  }
  
  data$C[data$D==1] = NA
  
  return(list(data=data,
              fY_sca_simpop = fY_sca_simpop,
              p_sc_simpop = p_sc_simpop,
              p_c_simpop = p_c_simpop,
              p_s.c_simpop = p_s.c_simpop,
              p_c.s_simpop = p_c.s_simpop,
              F_Y_sa = F_Y_sa,
              F_Y_scc = F_Y_scc,
              F_diff_Y_sa = F_diff_Y_sa
  ))
  
}

#########################
## bootstrap functions ##
#########################

bounds.frechet.boot <- function(data_use, J, boot_num = 1000){
  
  ## bootstrapping
  bounds.boot.list = lapply(1:boot_num, function(j){
    # cat('bounds boot', j, '\n')
    tryCatch({
      bounds.frechet(data_use[sample.int(nrow(data_use), replace = TRUE),],
                     J = J,
                     effects.mat=effects.mat,
                     posterior = FALSE,
                     n_dir_MC = n_dir_MC,
                     n_norm_MC = n_norm_MC,
                     hpd = TRUE,
                     alpha = .95,
                     sensitivity = TRUE,
                     rhos = seq(0, 12, 1)
      )
    },
    error = function(e){
      NA
    })
  })
  
  ## process bootstrap results and convert to ci
  
  failed_bootstrap_draw = sapply(bounds.boot.list, length) == 1
  bounds.boot.list = bounds.boot.list[!failed_bootstrap_draw]
  
  ## naive strata
  naive_strata.boot.arr = sapply(
    bounds.boot.list,
    function(x){
      as.matrix(x$naive[['strata']])
    },
    simplify = 'array'
  )
  naive_strata.boot <- as.data.frame(apply(naive_strata.boot.arr, 1:2, quantile, .025))
  colnames(naive_strata.boot)[match('naive', colnames(naive_strata.boot))] <-
    'min_cilo'
  naive_strata.boot$max_cihi <- apply(naive_strata.boot.arr[,'naive',], 1, quantile, .975)
  
  ## naive effects
  naive_effects.boot.arr = sapply(
    bounds.boot.list,
    function(x){
      as.matrix(x$naive[['effects']])
    },
    simplify = 'array'
  )
  naive_effects.boot <- as.data.frame(apply(naive_effects.boot.arr, 1:2, quantile, .025))
  colnames(naive_effects.boot)[match('naive', colnames(naive_effects.boot))] <-
    'min_cilo'
  naive_effects.boot$max_cihi <- apply(naive_effects.boot.arr[,'naive',], 1, quantile, .975)
  
  ## bounds on strata
  strata.boot.arr = sapply(
    bounds.boot.list,
    function(x){
      as.matrix(x[['strata']])
    },
    simplify = 'array'
  )
  strata.boot <- as.data.frame(apply(strata.boot.arr, 1:2, quantile, .025))
  strata.boot$max <- apply(strata.boot.arr[,'max',], 1, quantile, .975)
  colnames(strata.boot)[match(c('min', 'max'), colnames(strata.boot))] <-
    c('min_cilo', 'max_cihi')
  
  ## bounds on effects
  effects.boot.arr = sapply(
    bounds.boot.list,
    function(x){
      as.matrix(x[['effects']])
    },
    simplify = 'array'
  )
  effects.boot <- as.data.frame(apply(effects.boot.arr, 1:2, quantile, .025))
  effects.boot$max <- apply(effects.boot.arr[,'max',], 1, quantile, .975)
  colnames(effects.boot)[match(c('min', 'max'), colnames(effects.boot))] <-
    c('min_cilo', 'max_cihi')
  
  ## NEW Estimate 
  bounds.boot.new_13_m1 <- bounds.boot.new_23_m2 <- matrix(NA, nrow = length(bounds.boot.list), ncol = 2)
  for(x in 1:length(bounds.boot.list)){
    if(any(is.na(bounds.boot.list[[x]])) == FALSE){
      
      prop_c1 <- bounds.boot.list[[x]]$prop_c[-1]/sum(bounds.boot.list[[x]]$prop_c[-1])
      prop_c2 <- bounds.boot.list[[x]]$prop_c[-2]/sum(bounds.boot.list[[x]]$prop_c[-2])
      
      bounds.boot.new_13_m1[x, 1:2] <- c(t(bounds.boot.list[[x]]$effects[c(5, 6), -c(1,2,3), x]) %*% 
                                           prop_c1)
      bounds.boot.new_23_m2[x, 1:2] <- c(t(bounds.boot.list[[x]]$effects[c(7, 9), -c(1,2,3), x]) %*% 
                                           prop_c2)
      
      # bounds.boot.new_13_m1[x, 1:2] <- c(t(bounds.boot.list[[x]]$effects[c(5, 6), -c(1,2,3), x]) %*% 
      #                                      bounds.boot.list[[x]]$prop_c[-1])
      # bounds.boot.new_23_m2[x, 1:2] <- c(t(bounds.boot.list[[x]]$effects[c(7, 9), -c(1,2,3), x]) %*% 
      #                                      bounds.boot.list[[x]]$prop_c[-2])
    }
  }
  effect_13_m1 <- c(apply(bounds.boot.new_13_m1[,c(1,2)], 2, mean, na.rm = TRUE),
                    quantile(bounds.boot.new_13_m1[,1], prob = 0.025, na.rm = TRUE),
                    quantile(bounds.boot.new_13_m1[,2], prob = 0.975, na.rm = TRUE))
  effect_23_m2 <- c(apply(bounds.boot.new_23_m2[,c(1,2)], 2, mean, na.rm = TRUE),
                    quantile(bounds.boot.new_23_m2[,1], prob = 0.025, na.rm = TRUE),
                    quantile(bounds.boot.new_23_m2[,2], prob = 0.975, na.rm = TRUE))
  effect_new <- rbind(effect_13_m1, effect_23_m2)
  rownames(effect_new) <- c("13_m1", "23_m2")
  colnames(effect_new) <- c("min", "max", "min_cilo", "max_cihi")
  
  ## sensitivity on strata
  sens_strata.boot.arr = sapply(
    bounds.boot.list,
    function(x){
      as.matrix(x$sensitivity[['strata']])
    },
    simplify = 'array'
  )
  sens_strata.boot <-
    as.data.frame(apply(sens_strata.boot.arr, 1:2, quantile, .025, na.rm = TRUE))
  sens_strata.boot$max <-
    apply(sens_strata.boot.arr[,'max',], 1, quantile, .975, na.rm = TRUE)
  colnames(sens_strata.boot)[match(c('min', 'max'), colnames(sens_strata.boot))] <-
    c('min_cilo', 'max_cihi')
  sens_strata.boot$failed_lp <-
    apply(sens_strata.boot.arr[,'min',], 1, function(x) mean(is.na(x)))
  
  ## sensitivity on effects
  sens_effects.boot.arr = sapply(
    bounds.boot.list,
    function(x){
      as.matrix(x$sensitivity[['effects']])
    },
    simplify = 'array'
  )
  sens_effects.boot <-
    as.data.frame(apply(sens_effects.boot.arr, 1:2, quantile, .025, na.rm = TRUE))
  sens_effects.boot$max <-
    apply(sens_effects.boot.arr[,'max',], 1, quantile, .975, na.rm = TRUE)
  colnames(sens_effects.boot)[match(c('min', 'max'), colnames(sens_effects.boot))] <-
    c('min_cilo', 'max_cihi')
  sens_effects.boot$failed_lp <-
    apply(sens_effects.boot.arr[,'min',], 1, function(x) mean(is.na(x)))
  
  return(list(
    naive_strata.ci = naive_strata.boot,
    naive_effects.ci = naive_effects.boot,
    strata.ci = strata.boot,
    effects.ci = effects.boot,
    sens_strata.ci = sens_strata.boot,
    sens_effects.ci = sens_effects.boot,
    failed_boot = mean(failed_bootstrap_draw),
    effect_new  = effect_new
  ))
  
}